/*
 * @file arma_bsxfun.h
 *
 * Broadcast computation based on Armadillo
 *
 * @author Dahua Lin
 */

#ifndef _MSC_VER
#pragma once
#endif

#ifndef GMLPP_ARMA_BSXFUN_H_
#define GMLPP_ARMA_BSXFUN_H_

#include "arma_fun.h"

namespace gml
{

	/******************************************************
	 *
	 *  Generic evaluation functions
	 *
	 ******************************************************/

	/**
	 * Broadcast expansion of rows
	 */

	template<class Fun>
	class glue_bsxr
	{
	public:
		glue_bsxr(Fun f) : fun(f)
		{
		}

		template<typename T1, typename T2, typename eTR>
		inline static void apply(Mat<eTR>& out, const arma::Glue<T1, T2, glue_bsxr>& X)
		{
			typedef typename Fun::result_type eT;

			const arma::unwrap<T1> uA(X.A);
			const arma::unwrap<T2> uB(X.B);

			direct_bsxr(uA.M, uB.M, out);
		}

	private:
		template<typename eT1, typename eT2, typename eTR>
		static void direct_bsxr(const Mat<eT1>& X, const Mat<eT2>& row, Mat<eTR>& R)
		{
			uword m = X.n_rows;
			uword n = X.n_cols;
			R.set_size(m, n);

			const eT1 *x = X.memptr();
			const eT2 *y = row.memptr();
			eTR *r = R.memptr();

			for (uword j = 0; j < n; ++j, x += m, r += m)
			{
				eT2 yv = y[j];
				for (uword i = 0; i < m; ++i)
					r[i] = fun(x[i], yv);
			}
		}

	private:
		Fun fun;
	};

	template<class Fun, typename T1, typename T2>
	inline const GlueEx<T1, T2, glue_bsxr<Fun> >
	bsxr(Fun fun, const arma::Base<typename T1::elem_type,T1>& A, const arma::Base<typename T1::elem_type,T2>& B)
	{
		return GlueEx<T1, T2, glue_bsxr<Fun> >(glue_bsxr<Fun>(fun), A.get_ref(), B.get_ref());
	}


/*
	template<class Fun, typename T1, typename T2, typename TR>
	void bsxr(Fun fun, const Mat<T1>& X, const Mat<T2>& row, Mat<TR>& R)
	{
		uword m = X.n_rows;
		uword n = X.n_cols;

		check_arg(row.n_rows == 1 && row.n_cols == n && R.n_rows == m && R.n_cols == n,
				"dimension mismatch for bsxr");

		const T1 *x = X.memptr();
		const T2 *y = row.memptr();
		TR *r = R.memptr();

		for (uword j = 0; j < n; ++j, x += m, r += m)
		{
			T2 yv = y[j];
			for (uword i = 0; i < m; ++i)
				r[i] = fun(x[i], yv);
		}
	}

	template<class Fun, typename T1, typename T2>
	inline Mat<typename Fun::result_type> bsxr(Fun fun, const Mat<T1>& X, const Mat<T2>& row)
	{
		Mat<typename Fun::result_type> R(X.n_rows, X.n_cols);
		bsxr(fun, X, row, R);
		return R;
	}
	*/

	template<typename T, typename TR>
	inline void bsxr_plus(const Mat<T>& X, const Mat<T>& row, Mat<TR>& R)
	{
		bsxr(std::plus<T>(), X, row, R);
	}

	template<typename T>
	inline Mat<T> bsxr_plus(const Mat<T>& X, const Mat<T>& row)
	{
		return bsxr(std::plus<T>(), X, row);
	}

	template<typename T, typename TR>
	inline void bsxr_minus(const Mat<T>& X, const Mat<T>& row, Mat<TR>& R)
	{
		bsxr(std::minus<T>(), X, row, R);
	}

	template<typename T>
	inline Mat<T> bsxr_minus(const Mat<T>& X, const Mat<T>& row)
	{
		return bsxr(std::minus<T>(), X, row);
	}

	template<typename T, typename TR>
	inline void bsxr_times(const Mat<T>& X, const Mat<T>& row, Mat<TR>& R)
	{
		bsxr(std::multiplies<T>(), X, row, R);
	}

	template<typename T>
	inline Mat<T> bsxr_times(const Mat<T>& X, const Mat<T>& row)
	{
		return bsxr(std::multiplies<T>(), X, row);
	}

	template<typename T, typename TR>
	inline void bsxr_divides(const Mat<T>& X, const Mat<T>& row, Mat<TR>& R)
	{
		bsxr(std::divides<T>(), X, row, R);
	}

	template<typename T>
	inline Mat<T> bsxr_divides(const Mat<T>& X, const Mat<T>& row)
	{
		return bsxr(std::divides<T>(), X, row);
	}


	/**
	 * Broadcast expansion of columns
	 */

	template<class Fun, typename T1, typename T2, typename TR>
	void bsxc(Fun fun, const Mat<T1>& X, const Mat<T2>& col, Mat<TR>& R)
	{
		uword m = X.n_rows;
		uword n = X.n_cols;
		check_arg(col.n_rows == m && col.n_cols == 1 && R.n_rows == m && R.n_cols == n,
				"dimension mismatch for bsxc");

		const T1 *x = X.memptr();
		const T2 *y = col.memptr();
		TR *r = R.memptr();

		for (uword j = 0; j < n; ++j, x += m, r += m)
		{
			for (uword i = 0; i < m; ++i)
				r[i] = fun(x[i], y[i]);
		}
	}

	template<class Fun, typename T1, typename T2>
	inline Mat<typename Fun::result_type> bsxc(Fun fun, const Mat<T1>& X, const Mat<T2>& col)
	{
		Mat<typename Fun::result_type> R(X.n_rows, X.n_cols);
		bsxc(fun, X, col, R);
		return R;
	}

	template<typename T, typename TR>
	inline void bsxc_plus(const Mat<T>& X, const Mat<T>& col, Mat<TR>& R)
	{
		bsxc(std::plus<T>(), X, col, R);
	}

	template<typename T>
	inline Mat<T> bsxc_plus(const Mat<T>& X, const Mat<T>& col)
	{
		return bsxc(std::plus<T>(), X, col);
	}

	template<typename T, typename TR>
	inline void bsxc_minus(const Mat<T>& X, const Mat<T>& col, Mat<TR>& R)
	{
		bsxc(std::minus<T>(), X, col, R);
	}

	template<typename T>
	inline Mat<T> bsxc_minus(const Mat<T>& X, const Mat<T>& col)
	{
		return bsxc(std::minus<T>(), X, col);
	}

	template<typename T, typename TR>
	inline void bsxc_times(const Mat<T>& X, const Mat<T>& col, Mat<TR>& R)
	{
		bsxc(std::multiplies<T>(), X, col, R);
	}

	template<typename T>
	inline Mat<T> bsxc_times(const Mat<T>& X, const Mat<T>& col)
	{
		return bsxc(std::multiplies<T>(), X, col);
	}

	template<typename T, typename TR>
	inline void bsxc_divides(const Mat<T>& X, const Mat<T>& col, Mat<TR>& R)
	{
		bsxc(std::divides<T>(), X, col, R);
	}

	template<typename T>
	inline Mat<T> bsxc_divides(const Mat<T>& X, const Mat<T>& col)
	{
		return bsxc(std::divides<T>(), X, col);
	}


	/**
	 * pairwise table
	 */
	template<class Fun, typename T1, typename T2, typename TR>
	void pairwise(Fun fun, const Mat<T1>& x, const Mat<T2>& y, Mat<TR>& R)
	{
		uword m = x.n_elem;
		uword n = y.n_elem;

		check_arg(x.is_vec() && y.is_vec() && R.n_rows == m && R.n_cols == n,
				"dimension mismatch for pairwise calculation");

		for (uword j = 0; j < n; ++j)
		{
			for (uword i = 0; i < m; ++i)
			{
				R.at(i, j) = fun(x[i], y[j]);
			}
		}
	}


	template<class Fun, typename T1, typename T2>
	Mat<typename Fun::result_type> pairwise(Fun fun, const Mat<T1>& x, const Mat<T2>& y)
	{
		Mat<typename Fun::result_type> R(x.n_elem, y.n_elem);
		pairwise(fun, x, y, R);
		return R;
	}

}

#endif /* ARMA_BSXFUN_H_ */
